// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/sparksql/tests/SparkFunctionBaseTest.h>

using namespace kumo::pollux::test;

namespace kumo::pollux::functions::sparksql::test {
namespace {

class ArrayConcatTest : public SparkFunctionBaseTest {
 protected:
  void testExpression(
      const std::string& expression,
      const std::vector<VectorPtr>& input,
      const VectorPtr& expected) {
    auto result = evaluate(expression, make_row_vector(input));
    assertEqualVectors(expected, result);
  }
};

// Concatenate two integer arrays.
TEST_F(ArrayConcatTest, intArray) {
  const auto intArray = make_array_vector<int64_t>(
      {{1, 2, 3, 4}, {3, 4, 5}, {7, 8, 9}, {10, 20, 30}});
  const auto emptyArray = make_array_vector<int64_t>({{9, 8}, {}, {777}, {}});

  VectorPtr expected = make_array_vector<int64_t>({
      {1, 2, 3, 4, 9, 8},
      {3, 4, 5},
      {7, 8, 9, 777},
      {10, 20, 30},
  });
  testExpression("concat(c0, c1)", {intArray, emptyArray}, expected);

  expected = make_array_vector<int64_t>({
      {9, 8, 1, 2, 3, 4},
      {3, 4, 5},
      {777, 7, 8, 9},
      {10, 20, 30},
  });
  testExpression("concat(c0, c1)", {emptyArray, intArray}, expected);
}

// Concatenate two integer arrays with null.
TEST_F(ArrayConcatTest, nullArray) {
  const auto intArray =
      make_array_vector<int64_t>({{1, 2, 3, 4}, {7, 8, 9}, {10, 20, 30}});
  const auto nullArray = make_nullable_array_vector<int64_t>({
      {{std::nullopt, std::nullopt}},
      std::nullopt,
      {{1}},
  });

  VectorPtr expected = make_nullable_array_vector<int64_t>({
      {{1, 2, 3, 4, std::nullopt, std::nullopt}},
      std::nullopt,
      {{10, 20, 30, 1}},
  });
  testExpression("concat(c0, c1)", {intArray, nullArray}, expected);

  expected = make_nullable_array_vector<int64_t>({
      {{std::nullopt, std::nullopt, 1, 2, 3, 4}},
      std::nullopt,
      {{1, 10, 20, 30}},
  });
  testExpression("concat(c0, c1)", {nullArray, intArray}, expected);
}

TEST_F(ArrayConcatTest, arity) {
  const auto array1 =
      make_array_vector<int64_t>({{1, 2, 3, 4}, {3, 4, 5}, {7, 8, 9}});
  const auto array2 = make_array_vector<int64_t>({{9, 8}, {}, {777}});
  const auto array3 = make_array_vector<int64_t>({{123, 42}, {55}, {10, 20, 30}});
  const auto array4 = make_nullable_array_vector<int64_t>({
      {{std::nullopt, std::nullopt}},
      std::nullopt,
      {{std::nullopt, std::nullopt, std::nullopt}},
  });

  testExpression("concat(c0)", {array1}, array1);

  VectorPtr expected = make_array_vector<int64_t>(
      {{1, 2, 3, 4, 9, 8, 123, 42}, {3, 4, 5, 55}, {7, 8, 9, 777, 10, 20, 30}});
  testExpression("concat(c0, c1, c2)", {array1, array2, array3}, expected);

  expected = make_array_vector<int64_t>(
      {{123, 42, 9, 8, 1, 2, 3, 4, 9, 8},
       {55, 3, 4, 5},
       {10, 20, 30, 777, 7, 8, 9, 777}});
  testExpression(
      "concat(c0, c1, c2, c3)", {array3, array2, array1, array2}, expected);

  expected = make_nullable_array_vector<int64_t>(
      {{{1, 2, 3, 4, std::nullopt, std::nullopt, 123, 42, 9, 8}},
       std::nullopt,
       {{7, 8, 9, std::nullopt, std::nullopt, std::nullopt, 10, 20, 30, 777}}});
  testExpression(
      "concat(c0, c1, c2, c3)", {array1, array4, array3, array2}, expected);
}

// Concatenate complex types.
TEST_F(ArrayConcatTest, complexTypes) {
  auto baseVector = make_array_vector<int64_t>(
      {{1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 5}, {6, 6}});

  // Create arrays of array vector using above base vector.
  // [[1, 1], [2, 2]]
  // [[3, 3], [4, 4]]
  // [[5, 5], [6, 6]]
  auto arrayOfArrays1 = make_array_vector({0, 2, 4}, baseVector);
  // [[1, 1], [2, 2], [3, 3]]
  // [[4, 4]]
  // [[5, 5], [6, 6]]
  auto arrayOfArrays2 = make_array_vector({0, 3, 4}, baseVector);

  // [[1, 1], [2, 2], [1, 1], [2, 2], [3, 3]]
  // [[3, 3], [4, 4], [4, 4]]
  // [[5, 5], [6, 6], [5, 5], [6, 6]]
  auto expected = make_array_vector(
      {0, 5, 8},
      make_array_vector<int64_t>(
          {{1, 1},
           {2, 2},
           {1, 1},
           {2, 2},
           {3, 3},
           {3, 3},
           {4, 4},
           {4, 4},
           {5, 5},
           {6, 6},
           {5, 5},
           {6, 6}}));

  testExpression("concat(c0, c1)", {arrayOfArrays1, arrayOfArrays2}, expected);
}

TEST_F(ArrayConcatTest, unknown) {
  const auto array1 = make_array_vector_from_json<UnknownValue>({"[null, null]"});
  const auto array2 = make_array_vector_from_json<UnknownValue>({"[null]"});

  const auto expected =
      make_array_vector_from_json<UnknownValue>({"[null, null, null]"});
  testExpression("concat(c0, c1)", {array1, array2}, expected);
}

} // namespace
} // namespace kumo::pollux::functions::sparksql::test
